import cv2
import face_recognition
import numpy as np
# 定义摄像头捕获的分辨率
FRAME_WIDTH = 640
FRAME_HEIGHT = 480

# 初始化摄像头
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)

# 定义人脸检测模型为 HOG，以提高速度
FACE_RECOGNITION_MODEL = 'hog'

# 定义上采样次数，减少这个值可以加快速度，但可能会降低小脸的检测率
NUMBER_OF_TIMES_TO_UPSAMPLE = 1

# 定义人脸编码的公差阈值
TOLERANCE = 0.6

# 已知人脸编码的列表（示例中为空，实际使用时应填入已知人脸的编码）
known_face_encodings = []

# 已知人名列表（示例中为空，实际使用时应填入对应的人名）
known_face_names = []

# 循环读取摄像头的每一帧
while True:
    # 读取一帧
    ret, frame = cap.read()
    if not ret:
        print("无法读取帧")
        break
    
    # 转换颜色空间从BGR到RGB
    rgb_frame = np.ascontiguousarray(frame[:, :, ::-1])
    
    # 检测帧中所有人脸的位置
    face_locations = face_recognition.face_locations(rgb_frame, model=FACE_RECOGNITION_MODEL, number_of_times_to_upsample=NUMBER_OF_TIMES_TO_UPSAMPLE)
    
    # 为每个人脸提取面部编码
    face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)
    
    # 遍历每个面部编码
    for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
        # 将当前面部编码与已知面部编码进行比较
        matches = face_recognition.compare_faces(known_face_encodings, face_encoding, TOLERANCE)
        
        # 确定面部编码是否匹配
        name = "Unknown"
        
        if True in matches:
            first_match_index = matches.index(True)
            name = known_face_names[first_match_index]
        
        # 在帧上绘制矩形框和名称
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)
        cv2.putText(frame, name, (left + 6, bottom - 6), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255), 1)
    
    # 显示帧
    cv2.imshow('Video', frame)
    
    # 按下'q'键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头资源
cap.release()
# 关闭所有OpenCV窗口
cv2.destroyAllWindows()